from typing import Callable, List, Optional, Set

from torch import Tensor

class Offloader:
    def __init__(self, filename: str, n_entries: int, backend: str = "aio") -> None: ...
    def async_write(self, tensor: Tensor, key: str, callback: Optional[Callable[[], None]] = None) -> None: ...
    def async_read(self, tensor: Tensor, key: str, callback: Optional[Callable[[], None]] = None) -> None: ...
    def sync_write(self, tensor: Tensor, key: str) -> None: ...
    def sync_read(self, tensor: Tensor, key: str) -> None: ...
    def sync_write_events(self) -> None: ...
    def sync_read_events(self) -> None: ...
    def synchronize(self) -> None: ...
    def async_writev(self, tensors: List[Tensor], key: str, callback: Optional[Callable[[], None]] = None) -> None: ...
    def async_readv(self, tensors: List[Tensor], key: str, callback: Optional[Callable[[], None]] = None) -> None: ...
    def sync_writev(self, tensors: List[Tensor], key: str) -> None: ...
    def sync_readv(self, tensors: List[Tensor], key: str) -> None: ...

def get_backends() -> Set[str]: ...
def probe_backend(backend: str) -> bool: ...

class AsyncFileWriter:
    def __init__(self, fd: int, n_entries: int, backend: str = "aio") -> None: ...
    def write(self, buffer: int, n_bytes: int, offset: int, callback: Optional[Callable[[], None]] = None) -> None: ...
    def write_tensor(self, tensor: Tensor, offset: int, callback: Optional[Callable[[], None]] = None, pinned: Optional[Tensor] = None) -> None: ...
    def synchronize(self) -> None: ...
